# -*- coding: utf-8 -*-
'''
Created on 2017年3月20日

@author: ZhuJiahui
'''

import os
import time
from file_utils.file_writer import quick_write_1d_to_text
from global_info.global_nlp import GlobalNLP

def text_extraction(read_directory, write_directory):
    
    dir_list = os.listdir(read_directory)
    separater_string = GlobalNLP.EN_COLUMN_DELIMITER
    
    for i in range(len(dir_list)):
        
        inner_read_directory = read_directory + '/' + dir_list[i]
        file_list = os.listdir(inner_read_directory)
        to_write = []
        
        for j in range(len(file_list)):
            to_write1 = []
            to_write2 = []
            
            content_flag = 0
            
            with open(inner_read_directory + '/' + file_list[j], 'rb') as f:
                for each_line in f:
                    try:
                        each_line_str = each_line.decode()
                    except UnicodeDecodeError:
                        each_line_str = ""
                    
                    if len(each_line_str.strip()) < 2:
                        content_flag = 1
                    
                    elif each_line_str.startswith('From:'):
                        to_write1.append(each_line_str.strip()[6:])
                
                    elif each_line_str.startswith('Subject:'):
                        to_write1.append(each_line_str.strip()[9:])
                    
                    elif each_line_str.startswith('Organization:'):
                        to_write1.append(each_line_str.strip()[14:])
                
                    elif each_line_str.startswith('Summary:'):
                        to_write1.append(each_line_str.strip()[9:])
                
                    elif each_line_str.startswith('Keywords:'):
                        to_write1.append(each_line_str.strip()[10:])
                
                    else:
                        if content_flag == 1:
                            refined_text = each_line_str.strip().replace('>', '')
                            refined_text = refined_text.replace('<', '')
                            to_write2.append(refined_text)
            
            if (len(' '.join(to_write2).strip()) > 10):        
                to_write.append(str(file_list[j]) + separater_string + ' '.join(to_write1) + separater_string + ' '.join(to_write2))
        
        quick_write_1d_to_text(write_directory + '/' + dir_list[i] + '.txt', to_write)

if __name__ == '__main__':
    
    start = time.clock()    
    now_directory = os.getcwd()
    root_directory = os.path.dirname(now_directory) + '/'
    read_directory1 = root_directory + u'dataset/20newsgroup/20news-bydate-train'
    write_directory1 = root_directory + u'dataset/20newsgroup/train_text'
    read_directory2 = root_directory + u'dataset/20newsgroup/20news-bydate-test'
    write_directory2 = root_directory + u'dataset/20newsgroup/test_text'
    
    if (not(os.path.exists(write_directory1))):
        os.mkdir(write_directory1)
    if (not(os.path.exists(write_directory2))):
        os.mkdir(write_directory2)
    
    text_extraction(read_directory1, write_directory1)
    text_extraction(read_directory2, write_directory2)

    print('Total time %f seconds' % (time.clock() - start))
